import torch
import torch.nn as nn

"""
包含lstm中的遗忘门和输入门，用以不同尺度的图像融合
"""


# 定义输入门
class Lstm(nn.Module):
    def __init__(self, input_size, hidden_size,  cell_size):
        """
        通过输入状态和隐藏状态来更新细胞状态，本项目中细胞状态代表实际的位姿预测
        :param input_size: 输入
        :param hidden_size: 隐藏状态
        :param cell_size: 细胞状态
        """
        super(Lstm, self).__init__()
        # 定义遗忘门中的线性层，用于遗忘细胞中的部分状态
        self.forget_layer = nn.Linear(input_size + hidden_size, cell_size)
        # 定义输入门中的两个线性层，一个用于sigmoid，一个用于tanh
        self.sigmoid_layer = nn.Linear(input_size + hidden_size, cell_size)
        self.tanh_layer = nn.Linear(input_size + hidden_size, cell_size)

    def forward(self, input_tensor, hidden_state, cell_state):
        # 将当前输入和上一个时间步的隐藏状态拼接起来
        combined = torch.cat((input_tensor, hidden_state), 1)

        # 通过遗忘门控制信息的保留与删减
        forget_gate = torch.sigmoid(self.forget_layer(combined))

        # 通过sigmoid层计算输入门的权重
        input_gate_sigmoid = torch.sigmoid(self.sigmoid_layer(combined))

        # 通过tanh层计算新的候选值
        input_gate_tanh = torch.tanh(self.tanh_layer(combined))

        # 细胞状态会与遗忘门相乘，输入门相加
        cell_state = forget_gate*cell_state
        output = input_gate_sigmoid*input_gate_tanh
        output = output + cell_state

        # 返回融合更新后的output
        return output
#     # 初始化输入门
#
#
# input_size = 5  # 输入特征维度
# hidden_size = 20  # 隐藏层维度
# input_gate = InputGate(input_size, hidden_size)
#
# # 创建一个随机输入和隐藏状态
# batch_size = 1  # 批次大小
# sequence_length = 10  # 序列长度
# input_tensor = torch.randn(batch_size, sequence_length, input_size)
# hidden_state = torch.randn(batch_size, hidden_size)
#
# # 前向传播，计算输入门
# input_gate_sigmoid, input_gate_tanh = input_gate(input_tensor[:, -1, :], hidden_state)  # 只取最后一个时间步的输入
#
# print("Input Gate Sigmoid Output:", input_gate_sigmoid)
# print("Input Gate Tanh Output:", input_gate_tanh)
